from django.core.management.base import BaseCommand
from django.utils import timezone
from tushare_data.tushare_api import TushareAPI
from tushare_data.constants import START_KDATA_DATE
import datetime as dt
import logging

logger = logging.getLogger(__name__)

class Command(BaseCommand):
    help = '从Tushare更新股票数据'

    def add_arguments(self, parser):
        parser.add_argument(
            '--start-date',
            type=str,
            help='开始日期 (YYYY-MM-DD格式)',
            required=False
        )
        parser.add_argument(
            '--end-date',
            type=str,
            help='结束日期 (YYYY-MM-DD格式)',
            required=False
        )
        parser.add_argument(
            '--type',
            type=str,
            choices=['all', 'basic', 'calendar', 'daily'],
            default='all',
            help='要更新的数据类型'
        )

    def handle(self, *args, **options):
        try:
            # 初始化 TushareAPI
            api = TushareAPI()
            
            # 处理日期参数
            start_date = options.get('start_date')
            if not start_date:
                start_date = START_KDATA_DATE
                
            end_date = options.get('end_date')
            if not end_date:
                end_date = dt.datetime.now().strftime('%Y-%m-%d')

            update_type = options.get('type')
            
            self.stdout.write(f"开始更新数据，类型: {update_type}, 时间范围: {start_date} 到 {end_date}")

            if update_type == 'all':
                results = api.update_all_data(start_date, end_date)
                self.stdout.write(self.style.SUCCESS(
                    f"成功更新所有数据:\n"
                    f"- 股票基本信息: {results.get('stock_basic', 0)}条记录\n"
                    f"- 交易日历: {results.get('trade_cal', 0)}条记录\n"
                    f"- 日线数据: {results.get('daily', 0)}条记录"
                ))
            
            elif update_type == 'basic':
                count = api.update_stock_basic()
                self.stdout.write(self.style.SUCCESS(f"成功更新股票基本信息: {count}条记录"))
            
            elif update_type == 'calendar':
                count = api.update_trade_calendar(
                    start_date=start_date.replace('-', ''),
                    end_date=end_date.replace('-', '')
                )
                self.stdout.write(self.style.SUCCESS(f"成功更新交易日历: {count}条记录"))
            
            elif update_type == 'daily':
                count = api.batch_update_daily_data(
                    start_date=start_date.replace('-', ''),
                    end_date=end_date.replace('-', '')
                )
                self.stdout.write(self.style.SUCCESS(f"成功更新日线数据: {count}条记录"))

        except Exception as e:
            self.stdout.write(self.style.ERROR(f"更新数据时发生错误: {str(e)}"))
            logger.error(f"更新数据时发生错误: {str(e)}")
            raise